import torch
import numpy as np
from kan import MultKAN
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error, max_error
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

def train_and_eval_kan(data, epochs=1000, hidden=10, lr=0.01, device='cpu', test_size=0.2, random_state=42):
    device = torch.device(device)
    X = data[['freq_all', 'flow', 'dt']].values.astype(np.float32)
    y = data['heat_load'].values.astype(np.float32).reshape(-1, 1)
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=random_state)
    X_train_tensor = torch.tensor(X_train, device=device)
    y_train_tensor = torch.tensor(y_train, device=device)
    X_test_tensor = torch.tensor(X_test, device=device)
    y_test_tensor = torch.tensor(y_test, device=device)

    model = MultKAN(width=[3, hidden, 1])
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = torch.nn.MSELoss()

    loss_list = []
    for epoch in range(epochs):
        optimizer.zero_grad()
        y_pred = model(X_train_tensor)
        loss = loss_fn(y_pred, y_train_tensor)
        loss.backward()
        optimizer.step()
        loss_list.append(loss.item())
        if epoch % 100 == 0:
            print(f'KAN Epoch {epoch}, Loss: {loss.item()}, LR: {optimizer.param_groups[0]["lr"]}')

    # 可视化loss曲线
    plt.figure()
    plt.plot(loss_list, label='Train Loss')
    plt.xlabel('Epoch')
    plt.ylabel('MSE Loss')
    plt.title('KAN训练Loss曲线')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()

    with torch.no_grad():
        y_pred_train = model(X_train_tensor).cpu().numpy().flatten()
        y_pred_test = model(X_test_tensor).cpu().numpy().flatten()
    y_train_true = y_train_tensor.cpu().numpy().flatten()
    y_test_true = y_test_tensor.cpu().numpy().flatten()

    print('\n【team-daniel/KAN回归结果】')
    print('【训练集】R2:', r2_score(y_train_true, y_pred_train))
    print('【训练集】MSE:', mean_squared_error(y_train_true, y_pred_train))
    print('【测试集】R2:', r2_score(y_test_true, y_pred_test))
    print('【测试集】MSE:', mean_squared_error(y_test_true, y_pred_test))
    print('【测试集】MAE:', mean_absolute_error(y_test_true, y_pred_test))
    print('【测试集】Max Error:', max_error(y_test_true, y_pred_test))


    # 假设 model 是 MultKAN 实例
    try:
        formula = model.symbolic_formula(var=['freq_all', 'flow', 'dt'])
        print('KAN 公式表达:')
        print(formula)
    except Exception as e:
        print('当前 KAN 包不支持自动导出公式:', e) 


    return model, y_pred_test 